package com.qying.rpc.registry;

import com.qying.rpc.common.RpcServiceHelper;
import com.qying.rpc.common.ServiceMeta;
import com.qying.rpc.registry.loadbalancer.ZKConsistenHashLoadBalancer;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
import java.io.IOException;
import java.util.List;
import java.util.Objects;

public class ZookeeperRegistryService implements RegistryService{

    public static final Integer BASE_SLEEP_TIME_MS = 1000;

    public static final Integer MAX_RETRIES = 3;

    private static final String ZK_BASE_PATH = "/qying_rpc";

    private final ServiceDiscovery<ServiceMeta> serviceDiscovery;

    /**
     * 注册地址
     */
    private String registryAddr;

    public ZookeeperRegistryService(String registryAddr) throws Exception {
        CuratorFramework client = CuratorFrameworkFactory.newClient(registryAddr, new ExponentialBackoffRetry(BASE_SLEEP_TIME_MS, MAX_RETRIES));
        client.start();
        JsonInstanceSerializer<ServiceMeta> serviceMetaJsonInstanceSerializer = new JsonInstanceSerializer<>(ServiceMeta.class);

        this.serviceDiscovery = ServiceDiscoveryBuilder.builder(ServiceMeta.class)
                .client(client)
                .serializer(serviceMetaJsonInstanceSerializer)
                .basePath(ZK_BASE_PATH)
                .build();

       this.serviceDiscovery.start();
    }

    @Override
    public void register(ServiceMeta serviceMeta) throws Exception {
        ServiceInstance<ServiceMeta> serviceInstance = ServiceInstance
                .<ServiceMeta>builder()
                .name(RpcServiceHelper.buildServiceKey(serviceMeta.getServiceName(),serviceMeta.getServiceVersion()))
                .address(serviceMeta.getServiceAddr())
                .port(serviceMeta.getServicePort())
                .payload(serviceMeta)
                .build();
        serviceDiscovery.registerService(serviceInstance);
    }

    @Override
    public void unRegister(ServiceMeta serviceMeta) throws Exception {
        ServiceInstance<ServiceMeta> serviceInstance = ServiceInstance
                .<ServiceMeta>builder()
                .name(serviceMeta.getServiceName())
                .address(serviceMeta.getServiceAddr())
                .port(serviceMeta.getServicePort())
                .payload(serviceMeta)
                .build();

        serviceDiscovery.unregisterService(serviceInstance);
    }

    @Override
    public ServiceMeta discovery(String serviceName, int invokerHashCode) throws Exception {
        List<ServiceInstance<ServiceMeta>> serviceInstances = (List<ServiceInstance<ServiceMeta>>) serviceDiscovery.queryForInstances(serviceName);

        ServiceInstance<ServiceMeta> instance = new ZKConsistenHashLoadBalancer().select(serviceInstances, invokerHashCode);

        if (Objects.nonNull(instance)) {
            return instance.getPayload();
        }

        return null;
    }

    @Override
    public void destory() throws IOException {
        serviceDiscovery.close();
    }
}
